import time, sys
import os
import matplotlib.pyplot as plt
import cv2 as cv
import pickle
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split


label = ['cl', 'dlrb', 'lyf', 'pyy', 'xlz', 'zyx']  # 成龙，迪丽热巴，刘亦菲，彭于晏，小李子，张艺兴

def createxy():
    # 导入照片 创建X和y
    outpath='X'
    if not os.path.exists(outpath):
        n = []
        Y = []
        for f in os.listdir('./img'): # 采用相对路径
            if (".jpg" not in f) or ('pre' in f):  # 如果不是图片，则跳过  predict预测图片暂时不用（含多个人）
                continue
            # 识别照片中的人脸
            img = cv.imread(f)
            img_gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
            img_gray = cv.equalizeHist(img_gray)  # 把图像的范围拉开

            # 1. 创建级联分类器
            face_cascade = cv.CascadeClassifier()
            # 2. 引入训练好的可用于人脸识别的级联分类器模型
            face_cascade.load("haarcascade_frontalface_alt.xml")
            # 3. 用此级联分类器识别图像中的所有人脸信息，返回一个包含有所有识别的人联系系的列表
            # 列表中每一个元素包含四个值：面部左上角的坐标(x,y) 以及面部的宽和高(w,h)
            faces = face_cascade.detectMultiScale(img_gray)

            for (x, y, w, h) in faces:
                dd = img[y:y + h, x:x + w]
                new_img = cv.resize(dd, (500, 500), interpolation=cv.INTER_NEAREST)  # 改变人脸大小图片
                new_img_ravel = new_img.ravel()
                n.append(new_img_ravel)
                # 创建y
                if 'cl' in f:
                    Y.append(0)
                elif 'dlrb' in f:
                    Y.append(1)
                elif 'lyf' in f:
                    Y.append(2)
                elif 'pyy' in f:
                    Y.append(3)
                elif 'xlz' in f:
                    Y.append(4)
                else:
                    Y.append(5)
        y = np.array(Y)
        X = np.array(n)
        # 直接把 x y 以二进制的形式 写到文件: X y中
        with open("y", 'wb') as p:
            pickle.dump(y, p)
            print('y创建成功')
        with open("X", 'wb') as p:
            pickle.dump(X, p)
            print('X创建成功')

    else:
        print("已存在")

    # 以二进制的形式把 X 中保存的对象读入
    with open('X', 'rb') as p:
        x=pickle.load(p)
    # 以二进制的形式把 y 中保存的对象读入
    with open('y', 'rb') as p:
        y=pickle.load(p)
    #  train_test_split
    print('x.shape',x.shape)

    X_train, X_test, y_train, y_test = train_test_split(x, y, test_size=0.25)
    print('X_train.shape', X_train.shape)
    return X_train, X_test, y_train, y_test

def get_faces(img):
    getimg = []
    img_gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    img_gray = cv.equalizeHist(img_gray)  # 对图像进行直方图增强

    # 1. 创建级联分类器
    face_cascade = cv.CascadeClassifier()
    # 2. 引入训练好的可用于人脸识别的级联分类器模型
    face_cascade.load("haarcascade_frontalface_alt.xml")
    # 3. 用此级联分类器识别图像中的所有人脸信息，返回一个包含有所有识别的人联系系的列表
    # 列表中每一个元素包含四个值：面部左上角的坐标(x,y) 以及面部的宽和高(w,h)
    faces = face_cascade.detectMultiScale(img_gray)

    # 4. 为图像中的所有面部画框
    for (x, y, w, h) in faces:
        cv.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)
        get = img[y:y + h, x:x + w]
        img = cv.resize(get, (500, 500), interpolation=cv.INTER_NEAREST)  # 变成244*244大小图片
        ravel_img = img.ravel()
        getimg.append(ravel_img)
    getface = np.array(getimg)
    return getface

def Ran(X_train, X_test, y_train, y_test):
    clf = RandomForestClassifier()
    clf.fit(X_train, y_train)
    acc = clf.score(X_test, y_test)
    print(acc)
    with open("acc", 'wb') as f:
        pickle.dump(acc, f)
    with open("clf", 'wb') as f:
        pickle.dump(clf, f)

if __name__ == '__main__':
    X_train, X_test, y_train, y_test = createxy()
    Ran(X_train, X_test, y_train, y_test)




